# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
"""WebSocket server exposing one or more YOLOv5 models."""

import argparse
import io
import json
import base64
import asyncio
import websockets
import torch
from PIL import Image

models = {}

async def handle_connection(websocket, path):
    """Handle a WebSocket connection."""
    try:
        async for message in websocket:
            # Parse incoming message (JSON format)
            data = json.loads(message)

            if "image" in data:
                # Decode the image from base64
                image_data = data["image"]
                im_bytes = io.BytesIO(base64.b64decode(image_data))
                im = Image.open(im_bytes)

                # Perform inference
                model = data.get("model", "yolov5s")
                if model in models:
                    results = models[model](im, size=640)  # reduce size=320 for faster inference
                    result_json = results.pandas().xyxy[0].to_json(orient="records")
                    # Send results back to the client
                    await websocket.send(json.dumps({"results": result_json}))
                else:
                    await websocket.send(json.dumps({"error": "Model not found"}))
            else:
                await websocket.send(json.dumps({"error": "No image provided"}))
    except Exception as e:
        await websocket.send(json.dumps({"error": str(e)}))

async def run_server(port=5000):
    """Start the WebSocket server."""
    async with websockets.serve(handle_connection, "localhost", port):
        print(f"WebSocket server running on ws://localhost:{port}")
        await asyncio.Future()  # run forever

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="WebSocket server exposing YOLOv5 model")
    parser.add_argument("--port", default=5000, type=int, help="port number")
    parser.add_argument("--model", nargs="+", default=["yolov5s"], help="model(s) to run, i.e. --model yolov5n yolov5s")
    opt = parser.parse_args()

    for m in opt.model:
        models[m] = torch.hub.load("ultralytics/yolov5", "custom", path="D:\\yolo\\yolov5-master\\yolov5-master\\runs\\train\\exp5\\weights\\best.pt", skip_validation=True)

    # Start the WebSocket server
    asyncio.run(run_server(opt.port))
